#include "attitude_control_node_PX4.hpp"
#include <cmath>
#include <chrono>
#include <iomanip>
#include <sstream>

using namespace std::chrono_literals;

AttitudeControlNodePX4::AttitudeControlNodePX4() : Node("attitude_control_node_px4") {
    loadParameters();
    setupPublishersSubscribers();
    initializeDataLogging();
    
    // Set up parameter callback
    parameter_callback_handle_ = this->add_on_set_parameters_callback(
        std::bind(&AttitudeControlNodePX4::parametersCallback, this, std::placeholders::_1));
    
    last_dock_time_ = this->get_clock()->now();
    last_boat_time_ = this->get_clock()->now();
    
    RCLCPP_INFO(this->get_logger(), "PX4 Attitude Control Node initialized");
}

AttitudeControlNodePX4::~AttitudeControlNodePX4() {
    if (log_file_.is_open()) {
        log_file_.close();
    }
}

// Similar parameter callback as APM version
rcl_interfaces::msg::SetParametersResult AttitudeControlNodePX4::parametersCallback(
    const std::vector<rclcpp::Parameter> &parameters) {
    
    rcl_interfaces::msg::SetParametersResult result;
    result.successful = true;
    
    for (const auto &param : parameters) {
        if (param.get_name() == "control.attitude_angle") {
            control_params_.attitude_angle = param.as_double();
            RCLCPP_INFO(this->get_logger(), "Updated attitude_angle to %.2f", control_params_.attitude_angle);
        } else if (param.get_name() == "control.flag_control") {
            control_params_.flag_control = param.as_int();
            RCLCPP_INFO(this->get_logger(), "Updated flag_control to %d", control_params_.flag_control);
        }
        // Add other parameter updates as needed...
    }
    
    return result;
}

void AttitudeControlNodePX4::loadParameters() {
    // Similar to APM version but without PWM parameters
    this->declare_parameter("control.kp_boat_angle", 1.0);
    this->declare_parameter("control.attitude_angle", 0.0);
    this->declare_parameter("control.use_3d_angle", true);
    this->declare_parameter("control.flag_control", 0);
    this->declare_parameter("control.csv_path", std::string("boat_data"));
    this->declare_parameter("default_throttle_output", 0.5);
    
    this->get_parameter("control.kp_boat_angle", control_params_.kp_boat_angle);
    this->get_parameter("control.attitude_angle", control_params_.attitude_angle);
    this->get_parameter("control.use_3d_angle", control_params_.use_3d_angle);
    this->get_parameter("control.flag_control", control_params_.flag_control);
    this->get_parameter("control.csv_path", control_params_.csv_path);
    this->get_parameter("default_throttle_output", control_params_.default_throttle_output);
    
    // Load timeout configuration
    this->declare_parameter("timeout.dock_timeout", 5.0);
    this->declare_parameter("timeout.boat_timeout", 2.0);
    
    this->get_parameter("timeout.dock_timeout", timeout_config_.dock_timeout);
    this->get_parameter("timeout.boat_timeout", timeout_config_.boat_timeout);
    
    // Load PID parameters
    this->declare_parameter("yaw.kp", 1.5);
    this->declare_parameter("yaw.ki", 0.001);
    this->declare_parameter("yaw.kd", 0.015);
    this->declare_parameter("yaw.output_limit_min", -1.0);
    this->declare_parameter("yaw.output_limit_max", 1.0);
    this->declare_parameter("yaw.integral_limit", 2.0);
    
    PIDParams yaw_params;
    this->get_parameter("yaw.kp", yaw_params.kp);
    this->get_parameter("yaw.ki", yaw_params.ki);
    this->get_parameter("yaw.kd", yaw_params.kd);
    this->get_parameter("yaw.output_limit_min", yaw_params.output_limit_min);
    this->get_parameter("yaw.output_limit_max", yaw_params.output_limit_max);
    this->get_parameter("yaw.integral_limit", yaw_params.integral_limit);
    yaw_pid_.updateParams(yaw_params);
    
    RCLCPP_INFO(this->get_logger(), "Parameters loaded - Control mode: %d", control_params_.flag_control);
}

void AttitudeControlNodePX4::setupPublishersSubscribers() {
    // Subscribers - same as APM
    state_sub_ = this->create_subscription<mavros_msgs::msg::State>(
        "/mavros/state", 10, std::bind(&AttitudeControlNodePX4::stateCallback, this, std::placeholders::_1));
    
    pose_sub_ = this->create_subscription<geometry_msgs::msg::PoseStamped>(
        "/mavros/local_position/pose", 10, std::bind(&AttitudeControlNodePX4::poseCallback, this, std::placeholders::_1));
    
    global_sub_ = this->create_subscription<sensor_msgs::msg::NavSatFix>(
        "/mavros/global_position/global", 10, std::bind(&AttitudeControlNodePX4::globalCallback, this, std::placeholders::_1));
    
    velocity_sub_ = this->create_subscription<geometry_msgs::msg::TwistStamped>(
        "/mavros/local_position/velocity_body", 10, std::bind(&AttitudeControlNodePX4::velocityCallback, this, std::placeholders::_1));
    
    dock_angle_3d_sub_ = this->create_subscription<std_msgs::msg::Float64>(
        "/yolo/dock_angle_3D", 10, std::bind(&AttitudeControlNodePX4::dockAngle3DCallback, this, std::placeholders::_1));
    
    dock_angle_2d_sub_ = this->create_subscription<std_msgs::msg::Float64>(
        "/yolo/dock_angle_2D", 10, std::bind(&AttitudeControlNodePX4::dockAngle2DCallback, this, std::placeholders::_1));
    
    boat_angle_sub_ = this->create_subscription<std_msgs::msg::Float64>(
        "/yolo/boat_angle", 10, std::bind(&AttitudeControlNodePX4::boatAngleCallback, this, std::placeholders::_1));
    
    // Publisher - ActuatorControl for PX4
    actuator_pub_ = this->create_publisher<mavros_msgs::msg::ActuatorControl>("/mavros/actuator_control", 10);
    
    // Service clients
    set_mode_client_ = this->create_client<mavros_msgs::srv::SetMode>("/mavros/set_mode");
    arming_client_ = this->create_client<mavros_msgs::srv::CommandBool>("/mavros/cmd/arming");
    
    // Timers
    dock_timer_ = this->create_wall_timer(1s, std::bind(&AttitudeControlNodePX4::checkDockTimeout, this));
    boat_timer_ = this->create_wall_timer(1s, std::bind(&AttitudeControlNodePX4::checkBoatTimeout, this));
    control_timer_ = this->create_wall_timer(50ms, std::bind(&AttitudeControlNodePX4::sendActuatorControl, this));
}

// Implement callback functions (similar to APM version)
void AttitudeControlNodePX4::stateCallback(const mavros_msgs::msg::State::SharedPtr msg) {
    current_state_ = *msg;
}

void AttitudeControlNodePX4::poseCallback(const geometry_msgs::msg::PoseStamped::SharedPtr msg) {
    position_x_ = msg->pose.position.x;
    position_y_ = msg->pose.position.y;
    position_z_ = msg->pose.position.z;
    
    double roll, pitch, yaw;
    quaternionToEuler(msg->pose.orientation, roll, pitch, yaw);
    
    if (yaw < 0) yaw += 2 * M_PI;
    current_yaw_ = yaw;
    
    roll_angle_ = roll * 180.0 / M_PI;
    pitch_angle_ = pitch * 180.0 / M_PI;
    yaw_angle_ = yaw * 180.0 / M_PI;

    if (control_params_.flag_control == 0 && current_yaw_record_flag_ == true){
        current_yaw_record_ = yaw;
        current_yaw_record_flag_ = false;
    }
}

void AttitudeControlNodePX4::sendActuatorControl() {
    auto msg = mavros_msgs::msg::ActuatorControl();

    detectAndCompensateYawJump();
    
    if (dock_angle_ != -1 || control_params_.flag_control == 0) {
        // Normal control
        double dt = 0.01;
        yaw_output_ = -yaw_pid_.calculate(current_yaw_, desired_yaw_, dt);
        throttle_output_ = control_params_.default_throttle_output;
        
        msg.controls[2] = yaw_output_;
        msg.controls[3] = throttle_output_;
    } else {
        // Emergency/lost signal mode
        yaw_output_ = -1.0;
        throttle_output_ = 0.0;
        
        msg.controls[2] = yaw_output_;
        msg.controls[3] = throttle_output_;
    }
    
    // Set other controls to zero
    for (int i = 0; i < 8; ++i) {
        if (i != 2 && i != 3) msg.controls[i] = 0.0;
    }
    
    msg.group_mix = 0;
    actuator_pub_->publish(msg);

    previous_yaw_ = current_yaw_;
    first_yaw_update_ = false;
    
    // Update desired yaw and log data every 10 cycles
    if (update_counter_ % 10 == 0) {
        updateDesiredYaw();
        logData();
        update_counter_ = 0;
    }
    update_counter_++;
}

// Add other missing method implementations (similar to APM version)
void AttitudeControlNodePX4::updateDesiredYaw() {
    if (control_params_.flag_control == 1 && dock_angle_ != -1) {
        // Vision control mode
        double dock_angle_filtered = (std::abs(dock_angle_) < 2) ? 0 : dock_angle_;
        double boat_angle_filtered = (std::abs(boat_angle_) < 2) ? 0 : boat_angle_;
        
        desired_yaw_ = -1 * (dock_angle_filtered + control_params_.kp_boat_angle * boat_angle_filtered) 
                      * M_PI / 180.0 + current_yaw_;
    } else if (control_params_.flag_control == 0) {
        // Attitude control mode
        desired_yaw_ = control_params_.attitude_angle * M_PI / 180.0 + current_yaw_record_;
    }
}

void AttitudeControlNodePX4::detectAndCompensateYawJump() {
    if (first_yaw_update_) {
        return;
    }
    
    double current_yaw_deg = current_yaw_ * 180.0 / M_PI;
    double previous_yaw_deg = previous_yaw_ * 180.0 / M_PI;
    double yaw_diff = current_yaw_deg - previous_yaw_deg;
    
    const double JUMP_THRESHOLD = 180.0;
    
    if (yaw_diff < -JUMP_THRESHOLD) {
        desired_yaw_ -= 2 * M_PI;
        RCLCPP_INFO(this->get_logger(), "Detected yaw jump from ~360 to ~0, adjusted desired_yaw by -360 degrees");
    } else if (yaw_diff > JUMP_THRESHOLD) {
        desired_yaw_ += 2 * M_PI;
        RCLCPP_INFO(this->get_logger(), "Detected yaw jump from ~0 to ~360, adjusted desired_yaw by +360 degrees");
    }
}

// Add placeholder implementations for other methods
void AttitudeControlNodePX4::initializeDataLogging() {
    std::string log_path = "/home/ubuntu/ros2_boat_ws/" + control_params_.csv_path + ".csv";
    log_file_.open(log_path, std::ios::app);
    
    if (log_file_.is_open()) {
        log_file_ << "time,latitude,longitude,position_x,position_y,position_z,"
                  << "roll_angle,pitch_angle,yaw_angle,v_body_x,"
                  << "dock_angle,dock_angle_2D,dock_angle_3D,boat_angle,"
                  << "current_yaw,desired_yaw,yaw_output,throttle_output\n";
    } else {
        RCLCPP_ERROR(this->get_logger(), "Failed to open CSV file for logging: %s", log_path.c_str());
    }
}

void AttitudeControlNodePX4::globalCallback(const sensor_msgs::msg::NavSatFix::SharedPtr msg) {
    latitude_ = msg->latitude;
    longitude_ = msg->longitude;
}

void AttitudeControlNodePX4::velocityCallback(const geometry_msgs::msg::TwistStamped::SharedPtr msg) {
    v_body_x_ = msg->twist.linear.x;
}

void AttitudeControlNodePX4::dockAngle3DCallback(const std_msgs::msg::Float64::SharedPtr msg) {
    dock_angle_3d_ = msg->data;
    last_dock_time_ = this->get_clock()->now();
    
    if (control_params_.use_3d_angle) {
        dock_angle_ = dock_angle_3d_;
    }
}

void AttitudeControlNodePX4::dockAngle2DCallback(const std_msgs::msg::Float64::SharedPtr msg) {
    dock_angle_2d_ = msg->data;
    last_dock_time_ = this->get_clock()->now();
    
    if (!control_params_.use_3d_angle) {
        dock_angle_ = dock_angle_2d_;
    }
}

void AttitudeControlNodePX4::boatAngleCallback(const std_msgs::msg::Float64::SharedPtr msg) {
    boat_angle_ = msg->data;
    last_boat_time_ = this->get_clock()->now();
}

void AttitudeControlNodePX4::checkDockTimeout() {
    auto now = this->get_clock()->now();
    if ((now - last_dock_time_).seconds() > timeout_config_.dock_timeout) {
        dock_angle_ = -1;
        RCLCPP_WARN_THROTTLE(this->get_logger(), *this->get_clock(), 5000, "Dock angle timeout - no signal received");
    }
}

void AttitudeControlNodePX4::checkBoatTimeout() {
    auto now = this->get_clock()->now();
    if ((now - last_boat_time_).seconds() > timeout_config_.boat_timeout) {
        boat_angle_ = 0;
        RCLCPP_WARN_THROTTLE(this->get_logger(), *this->get_clock(), 5000, "Boat angle timeout - setting to 0");
    }
}

void AttitudeControlNodePX4::quaternionToEuler(const geometry_msgs::msg::Quaternion& q, 
                                           double& roll, double& pitch, double& yaw) {
    double sinr_cosp = 2.0 * (q.w * q.x + q.y * q.z);
    double cosr_cosp = 1.0 - 2.0 * (q.x * q.x + q.y * q.y);
    roll = atan2(sinr_cosp, cosr_cosp);

    double sinp = 2.0 * (q.w * q.y - q.z * q.x);
    if (std::abs(sinp) >= 1)
        pitch = std::copysign(M_PI / 2, sinp);
    else
        pitch = asin(sinp);

    double siny_cosp = 2.0 * (q.w * q.z + q.x * q.y);
    double cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z);
    yaw = atan2(siny_cosp, cosy_cosp);
}

void AttitudeControlNodePX4::setOffboardMode() {
    auto request = std::make_shared<mavros_msgs::srv::SetMode::Request>();
    request->custom_mode = "OFFBOARD";
    
    auto future = set_mode_client_->async_send_request(request);
    if (rclcpp::spin_until_future_complete(this->get_node_base_interface(), future) == 
        rclcpp::FutureReturnCode::SUCCESS) {
        RCLCPP_INFO(this->get_logger(), "Offboard mode set successfully");
    } else {
        RCLCPP_ERROR(this->get_logger(), "Failed to set Offboard mode");
    }
}

void AttitudeControlNodePX4::armDrone() {
    auto request = std::make_shared<mavros_msgs::srv::CommandBool::Request>();
    request->value = true;
    
    auto future = arming_client_->async_send_request(request);
    if (rclcpp::spin_until_future_complete(this->get_node_base_interface(), future) == 
        rclcpp::FutureReturnCode::SUCCESS) {
        RCLCPP_INFO(this->get_logger(), "Vehicle armed successfully");
    } else {
        RCLCPP_ERROR(this->get_logger(), "Failed to arm vehicle");
    }
}

void AttitudeControlNodePX4::logData() {
    if (!log_file_.is_open()) return;
    
    std::string timestamp = getCurrentTimeString();
    
    log_file_ << timestamp << ","
              << latitude_ << "," << longitude_ << ","
              << position_x_ << "," << position_y_ << "," << position_z_ << ","
              << roll_angle_ << "," << pitch_angle_ << "," << yaw_angle_ << ","
              << v_body_x_ << ","
              << dock_angle_ << "," << dock_angle_2d_ << "," << dock_angle_3d_ << ","
              << boat_angle_ << ","
              << (current_yaw_ * 180.0 / M_PI) << ","
              << (desired_yaw_ * 180.0 / M_PI) << ","
              << yaw_output_ << "," << throttle_output_ << "\n";
}

std::string AttitudeControlNodePX4::getCurrentTimeString() {
    auto now = std::chrono::system_clock::now();
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) % 1000;
    
    std::time_t time_t = std::chrono::system_clock::to_time_t(now);
    std::tm tm = *std::localtime(&time_t);
    
    std::ostringstream oss;
    oss << std::put_time(&tm, "%Y-%m-%d=%H:%M:%S") 
        << "." << std::setw(3) << std::setfill('0') << ms.count();
    
    return oss.str();
}

void AttitudeControlNodePX4::run() {
    // Wait for connection
    while (rclcpp::ok() && !current_state_.connected) {
        rclcpp::spin_some(this->get_node_base_interface());
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    
    // Send initial commands for 100 cycles
    int initial_cycles = 100;
    auto start_time = this->get_clock()->now();
    
    while (rclcpp::ok() && initial_cycles > 0) {
        sendActuatorControl();
        rclcpp::spin_some(this->get_node_base_interface());
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
        initial_cycles--;
    }
    
    // Set mode and arm
    setOffboardMode();
    armDrone();
    
    RCLCPP_INFO(this->get_logger(), "Starting main control loop");
    rclcpp::spin(shared_from_this());
}